# coding: utf-8
import numpy as np
import shutil
import os
import sys
import time
import json
from common import *
from utils.Logger import Logger

if not os.path.exists("./log/"):
    os.mkdir("./log/")
log = Logger("./log/{}_{}.log".format(__file__.split('/')[-1],
                                             time.strftime("%Y%m%d-%H%M%S"), time.localtime), level='debug').logger

class ca():
    def __init__(self, name, bbox):
        self.name = name
        self.bbox = bbox

source_dir = "~/dataset/coco"
source_dir = os.path.expanduser(source_dir)
output_dir = DATASET_ROOT_DIR + "/coco2014"

# mkdir
mk_common_dir(output_dir)

fp = os.path.join(source_dir, "annotations/instances_train2014.json")
f = open(fp)
j = json.load(f)
f.close()


dict = {}
anno = j['annotations']

for a in anno:
    category_id = a['category_id']
    if not category_id in coco_label:
        continue

    #
    name = coco_label[category_id]
    bbox = a['bbox']
    bbox[2] = bbox[2] + bbox[0]
    bbox[3] = bbox[3] + bbox[1]
    bbox = [int(i) for i in bbox]
    c = ca(name, bbox)

    # image id
    image_id = a['image_id']
    if image_id not in dict:
        dict[image_id] = [c]
    else:
        dict[image_id] += [c]
log.info("extrace category id over, total image num is: {}".format(len(dict)))
global_idx = 0
images = j['images']
for i in images:
    file_name = i['file_name']
    h = i['height']
    w = i['width']
    c = 3
    id = i['id']

    objects = []
    if id in dict:
        objects = dict[id]

    if len(objects) <= 0:
        continue

    object_names = []
    object_bboxes = []
    for obj in objects:
        object_names += [obj.name]
        object_bboxes += [obj.bbox]

    log.info("{} have object: {}".format(file_name, object_names))
    # write xml
    image_name = "coco2014_{}".format(global_idx)
    global_idx += 1
    write_xml(os.path.join(output_dir, "Annotations"), image_name, "coco2014", [w,h,c], object_bboxes, object_names)

    image_path = os.path.join(output_dir, "JPEGImages/{}.jpg".format(image_name))

    shutil.copy(os.path.join(source_dir, "train2014/{}".format(file_name)), image_path)

log.info("process done, whole number is: {}".format(global_idx))